import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor):
        # 计算均方根
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        x_normalized = x / rms
        return self.gamma * x_normalized
    

"""
    旋转位置编码（Rotary Position Embedding，RoPE）

    位置编码 = sin(position / 10000^(2i / d))

    旋转位置编码的关键思想是，利用 旋转矩阵 来引入位置信息，而不是直接加到输入中。
    具体来说，通过旋转向量来表示每个位置的编码。
    RoPE 将位置编码和输入特征做 元素级的乘法，通过旋转操作使得模型在不同位置之间的关系变得更为有意义。
"""
class RotaryPositionEmbedding(nn.Module):
    def __init__(self, dim: int, max_len: int = 512):
        super(RotaryPositionEmbedding, self).__init__()
        self.dim = dim
        self.max_len = max_len
        # 生成位置编码
        """
                   1.0
            ___________________

             10000 ^ (x / dim) 
        """
        self.position = torch.arange(0, max_len).float()
        self.inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
            x: (batch_size, n_heads, seq_len, dim)
        """
        # 计算位置编码的正弦和余弦
        sin_cos_table = torch.einsum("i,j->ij", self.position, self.inv_freq)
        sin_cos_table = torch.cat([sin_cos_table.sin(), sin_cos_table.cos()], dim=-1)
        
        # 限制位置编码长度
        sin_cos_table = sin_cos_table[:x.size(2), :]

        # 在维度0上扩展位置编码
        position_emb = sin_cos_table.unsqueeze(0).unsqueeze(0)

        # 将位置编码与输入相乘
        return x * position_emb


if __name__ == '__main__':
    x = torch.randn(1, 32, 10, 64)
    rotary_pos_emb = RotaryPositionEmbedding(dim=64, max_len=100)
    output = rotary_pos_emb(x)
    print(output.shape)